import akshare as ak
import pandas as pd
import time
from typing import Dict, List, Optional, Tuple, Set
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from db_config import DBConfig  # 使用您提供的配置类

# 配置参数
MAX_WORKERS = 8
BATCH_SIZE = 1000  # 批量插入改为 1000
RETRY_TIMES = 10
DELAY = 1

def get_market_info(stock_code: str) -> Tuple[str, str]:
    """获取交易所和市场类型信息"""
    code = str(stock_code)
    if code.startswith('8'):    # 北交所
        return ('SZ', 'bj_se')
    elif code.startswith('688'): # 科创板
        return ('SH', 'sh_star')
    elif code.startswith('300'): # 创业板
        return ('SZ', 'sz_gem')
    elif code.startswith(('6', '9')): # 沪市
        return ('SH', 'sh_main')
    else:  # 深市主板
        return ('SZ', 'sz_main')

def fetch_with_retry(func, retries=RETRY_TIMES, delay=DELAY, **kwargs):
    """带重试机制的请求函数"""
    for attempt in range(retries):
        try:
            return func(**kwargs)
        except Exception as e:
            if attempt == retries - 1:
                print(f"请求最终失败: {str(e)}")
                raise
            sleep_time = delay * (attempt + 1)
            print(f"请求失败({attempt+1}/{retries})，{sleep_time}秒后重试...")
            time.sleep(sleep_time)

def get_delisted_codes() -> Set[str]:
    """获取退市股票代码集合"""
    print("获取退市股票列表...")
    delisted_codes = set()
    
    try:
        # 上交所退市股票
        df_sh = fetch_with_retry(ak.stock_info_sh_delist)
        if not df_sh.empty:
            codes = df_sh['公司代码'].astype(str).str.zfill(6).tolist()
            delisted_codes.update(codes)
        
        # 深交所退市股票
        df_sz = fetch_with_retry(ak.stock_info_sz_delist)
        if not df_sz.empty:
            codes = df_sz['证券代码'].astype(str).str.zfill(6).tolist()
            delisted_codes.update(codes)
        
        print(f"获取到 {len(delisted_codes)} 条退市股票代码")
        return delisted_codes
    except Exception as e:
        print(f"获取退市股票失败，使用备用列表: {str(e)}")
        return {'600432', '600485', '600656', '000033', '000047'}

def fetch_stock_details(stock_codes: List[str]) -> Dict[str, Dict]:
    """获取股票详情(行业+上市日期)"""
    print(f"开始获取 {len(stock_codes)} 只股票详情...")
    details = {}
    
    def worker(code: str) -> Tuple[str, Optional[Dict]]:
        try:
            info = fetch_with_retry(ak.stock_individual_info_em, symbol=code)
            if isinstance(info, pd.DataFrame):
                records = info.to_dict('records')
                return (code, {
                    'industry': next((r['value'] for r in records if r['item'] == '行业'), None),
                    'list_date': next((r['value'] for r in records if r['item'] == '上市时间'), None)
                })
            return (code, None)
        except Exception as e:
            print(f"股票 {code} 详情获取失败: {str(e)}")
            return (code, None)
    
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(worker, code) for code in stock_codes]
        for future in as_completed(futures):
            code, detail = future.result()
            if detail:
                details[code] = detail
    
    print(f"成功获取 {len(details)}/{len(stock_codes)} 只股票详情")
    return details

def parse_listing_date(date_str: Optional[str]) -> Optional[str]:
    """标准化上市日期格式（支持多种格式）"""
    if not date_str:
        return None
    try:
        date_str = str(date_str).strip()
        # 处理 20230806 格式
        if len(date_str) == 8 and date_str.isdigit():
            dt = datetime.strptime(date_str, '%Y%m%d')
            return dt.strftime('%Y-%m-%d')
        # 处理 2023-08-06 或 2023/08/06 格式
        try:
            dt = pd.to_datetime(date_str)
            return dt.strftime('%Y-%m-%d')
        except Exception:
            return None
    except Exception:
        return None

def fetch_stock_data() -> List[Dict]:
    """主数据获取函数"""
    print("="*40 + " 开始获取股票数据 " + "="*40)
    start_time = time.time()
    
    try:
        # 1. 获取全量股票列表
        print("获取全量A股列表...")
        df = fetch_with_retry(ak.stock_zh_a_spot_em)
        df = df[['代码', '名称']].rename(columns={'代码': 'stock_code', '名称': 'stock_name'})
        df['stock_code'] = df['stock_code'].astype(str).str.zfill(6)
        
        # 2. 过滤退市股票
        delisted_codes = get_delisted_codes()
        df = df[~df['stock_code'].isin(delisted_codes)]
        print(f"过滤后剩余股票数量: {len(df)}")
        
        # 3. 获取行业和上市日期
        stock_codes = df['stock_code'].tolist()
        details = fetch_stock_details(stock_codes)
        
        # 4. 构建最终数据
        data = []
        for _, row in df.iterrows():
            code = row['stock_code']
            detail = details.get(code, {})
            exchange, market = get_market_info(code)
            
            data.append({
                'stock_code': code,
                'stock_name': row['stock_name'],
                'exchange_name': exchange,
                'market_type': market,
                'industry_type': detail.get('industry'),
                'listing_date': parse_listing_date(detail.get('list_date')),
                'is_stock_st': 1 if ('ST' in row['stock_name'] or '*ST' in row['stock_name']) else 0
            })
        
        print(f"数据处理完成，共获取 {len(data)} 条数据，耗时 {time.time()-start_time:.1f}秒")
        return data
    except Exception as e:
        print(f"主流程出错: {str(e)}")
        return []

def save_to_database(data: List[Dict]):
    """使用DBConfig.execute_sql保存数据"""
    if not data:
        print("没有有效数据可保存")
        return
    
    sql = """INSERT INTO stock_basic (
        stock_code, stock_name, exchange_name, market_type,
        industry_type, listing_date, is_stock_st
    ) VALUES (
        %(stock_code)s, %(stock_name)s, %(exchange_name)s, %(market_type)s,
        %(industry_type)s, %(listing_date)s, %(is_stock_st)s
    ) ON DUPLICATE KEY UPDATE
        stock_name = VALUES(stock_name),
        industry_type = VALUES(industry_type),
        listing_date = VALUES(listing_date),
        is_stock_st = VALUES(is_stock_st),
        update_time = CURRENT_TIMESTAMP"""
    
    try:
        # 分批处理
        total = 0
        for i in range(0, len(data), BATCH_SIZE):
            batch = data[i:i+BATCH_SIZE]
            try:
                affected_rows = DBConfig.execute_sql(sql, batch)
                total += affected_rows
                print(f"已提交 {min(i+BATCH_SIZE, len(data))}/{len(data)} 条，影响行数: {affected_rows}")
            except Exception as e:
                print(f"批次插入失败: {str(e)}")
                # 逐条重试
                for item in batch:
                    try:
                        DBConfig.execute_sql(sql, item)
                        total += 1
                    except Exception as e:
                        print(f"记录插入失败: {item.get('stock_code')}, 错误: {str(e)}")
        
        print(f"数据保存完成，成功插入/更新 {total} 条记录")
    except Exception as e:
        print(f"数据库操作失败: {str(e)}")

if __name__ == "__main__":
    try:
        start_time = time.time()
        stock_data = fetch_stock_data()
        if stock_data:
            save_to_database(stock_data)
        print(f"任务执行完毕，总耗时 {time.time()-start_time:.2f} 秒")
    except Exception as e:
        print(f"程序异常终止: {str(e)}")
